info_a = "A看到B和C的盒子内的球，并且说：“B和C两个盒子内的球不全是白色的。”"
info_b = "B看到A和C的盒子内的球，并且说：“A和C两个盒子内的球都是白色的。”"
info_c = "C看到A和B的盒子内的球，但未发表任何意见。"

info = f"""
有三个人A、B、C，每个人都拿着一个盒子。盒子内可能有一颗红球或者白球，
他们不能看到自己的盒子里球的颜色，而能看到其他人盒子内球的颜色。
我们知道以下信息，且三人均只说真话，不说假话：
{info_a}
{info_b}
{info_c}
我们需要推理出A、B、C三个人的盒子内的球的颜色。
"""

from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())

#from openai import OpenAI
#client = OpenAI()
from zhipuai import ZhipuAI
client = ZhipuAI()

def generate_text(prompt):
    messages = [{"role":"user", "content":prompt}]
    print(f"\n\nprompt=\n{prompt}")
    response = client.chat.completions.create(
        model="glm-4-plus",
        messages=messages
    )
    #print(f"response=\n{response}")
    resultContent = response.choices[0].message.content.strip()
    print(f"resultContent=\n{resultContent}")
    return resultContent

from anytree import Node, RenderTree
def create_logic_tree():
    # 创建根节点
    root = Node("推理盒子里的球的颜色")
    
    # 创建B的状态节点
    b_red = Node("B: 红", parent=root)
    b_white = Node("B: 白", parent=root)
    
    # 创建C的状态节点
    for b_node in [b_red, b_white]:
        Node("C: 红", parent=b_node)
        Node("C: 白", parent=b_node)
    
    # 创建A的状态节点
    for b_node in [b_red, b_white]:
        for c_node in b_node.children:
            Node("A: 红", parent=c_node)
            Node("A: 白", parent=c_node)
    
    return root

def prune_tree(state, condition):
    prompt = f"""判断条件'{condition}'和状态: '{state}'是否矛盾？
    没有矛盾用'no'回答，存在矛盾用'yes'回答，
    请仅判断该条件与状态是否矛盾，
    不要进行任何推理判断，也不要进行任何假设。
    简要说明判断依据。"""
    result = generate_text(prompt)
    return "no" in result.lower()

# 遍历树进行剪枝
def solve_question(root):
    states = []    
    for b_node in root.children:
        b_state = b_node.name
        for c_node in b_node.children:
                c_state = f"{b_state}, {c_node.name}"
                # 根据A的陈述进行剪枝
                if prune_tree(c_state, info_a):
                    for a_node in c_node.children:
                        a_state = f"{c_node.name}, {a_node.name}"
                        # 根据B的陈述进行剪枝
                        if prune_tree(a_state, info_b):
                            state = f"({a_node.name},{b_node.name},{c_node.name})"
                            states.append(state)    
    return states

# 打印树结构
def display_tree(root):
    for pre, fill, node in RenderTree(root):
        print("%s%s" % (pre, node.name))

def display_solutions(states):
    print(f"\n待求解问题:{info}")
    if states:
        print("\n逻辑推理的结果为:")
        for state in states:
            print(state)
    else:
        print("\n没有找到满足条件的解。")

def main():
    root = create_logic_tree()
    print("生成的思维树:")
    display_tree(root)
    
    solutions = solve_question(root)
    display_solutions(solutions)

if __name__ == "__main__":
    main()